# coding=utf-8
from __future__ import print_function, absolute_import
from gm.api import *
import datetime
import numpy as np
import pandas as pd
import statsmodels.api as sm

# 策略中必须有init方法
def init(context):
    # 每天的09:30 定时执行algo任务
    schedule(schedule_func=algo, date_rule='1d', time_rule='09:30:00') 
    context.num = 20
    context.alpha_STR_history_num = 20# STR因子的历史回溯天数
    # 大盘指数
    context.symbol = 'SHSE.000001'#'SHSE.000905'#


def algo(context):
    # 当前时间str
    now_date = context.now.strftime("%Y-%m-%d %H:%M:%S")
    # 下一个交易日
    next_date = get_N_trading_date(context,date=now_date,counts=1,model='next',exchange='SHSE')
    # 上一个交易日
    last_date = get_N_trading_date(context,date=now_date,counts=1,model='previous',exchange='SHSE')
    # 计算Dual Thrust通道
    # 每月最后一个交易日时换股
    if now_date[5:7]!=next_date[5:7]:
        # 获取全A股票（剔除停牌股和ST股）
        all_stocks,all_stocks_str = get_normal_stocks(context.now)
        # 计算STR因子
        STR = get_alpha_STR(context,all_stocks_str,last_date,context.alpha_STR_history_num)
        # 获取最小因子的前N只股票
        to_buy = list(STR.replace([-np.inf,np.inf],np.nan).dropna().sort_values(ascending=True)[:context.num].index)
        print(context.now,'待买入股票：{}'.format(to_buy))

        # 交易
        positions = context.account().positions()
        holding_symbols = [posi['symbol'] for posi in positions]
        new_prices = history(symbol=holding_symbols, frequency='1d', start_time=now_date, end_time=now_date, fields='symbol,close', adjust=ADJUST_PREV, adjust_end_time=context.backtest_end_time, df=False)
        # 平不在标的池的股票（注：日频市价单，backtest_match_mode=1，以当日收盘价进行交易）
        for position in positions:
            symbol = position['symbol']
            if symbol not in to_buy:
                new_price = list(filter(lambda x:x[r'symbol'] == symbol, new_prices))
                if len(new_price)==0:continue
                order_target_percent(symbol=symbol, percent=0, order_type=OrderType_Market,position_side=PositionSide_Long,price=new_price[0]['close']*0.99)

        # 获取股票的权重
        percent = 1 / len(to_buy)
        # 买在标的池中的股票（注：日频市价单，backtest_match_mode=1，以当日收盘价进行交易） 
        new_prices = history(symbol=to_buy, frequency='1d', start_time=now_date, end_time=now_date, fields='symbol,close', adjust=ADJUST_PREV, adjust_end_time=context.backtest_end_time, df=False)
        for symbol in to_buy:
            # 收盘价（日频数据）
            new_price = list(filter(lambda x:x[r'symbol'] == symbol, new_prices))
            if len(new_price)==0:continue
            order_target_percent(symbol=symbol, percent=percent, order_type=OrderType_Market,position_side=PositionSide_Long,price=new_price[0]['close']*1.01)


def get_alpha_STR(context,security,date,counts=1):
    """计算STR因子数据（包含date当日）
    :param security：目标股票，'***,***,***'格式
    :param date：目标日期
    :param counts：历史回溯天数
    """
    # 计算STR
    start_date = get_N_trading_date(context,date=date,counts=counts-1,model='previous',exchange='SHSE')
    turnrate = stk_get_daily_basic_pt_new(symbols=security, fields='turnrate', start_date=start_date, end_date=date, df=True)
    turnrate = turnrate.drop_duplicates(subset=['trade_date','symbol'],keep='first').set_index(['trade_date','symbol'])
    turnrate = turnrate[['turnrate']].unstack()
    turnrate.columns = turnrate.columns.droplevel(level=0)
    turnrate = turnrate.fillna(method='ffill').iloc[-counts:,:]
    STR = turnrate.std().dropna()
    # 去极值
    STR = winsorize_med(STR)
    # 标准化
    STR = standardlize(STR)
    # 中性化
    STR = neutralize_MarketValue(STR,date)
    return STR


def stk_get_daily_basic_pt_new(symbols, fields, start_date, end_date, df=True):
    data_df = pd.DataFrame()
    for trade_date in get_trading_dates(exchange='SZSE', start_date=start_date, end_date=end_date):
        data = stk_get_daily_basic_pt(symbols=symbols, fields=fields, trade_date=trade_date, df=df)
        data_df = pd.concat([data_df,data])
    return data_df


def winsorize_med(data, scale=3, inclusive=True, inf2nan=True):
    """
    去极值
    :param data：待处理数据[Series]
    :param scale：标准差倍数，默认为3
    :param inclusive：True为将边界外的数值调整为边界值，False为将边界外的数值调整为NaN
    :param inf2nan：True为将inf转化为nan，False不转化
    """
    data = data.astype('float')
    if inf2nan:
        data = data.replace([np.inf, -np.inf], np.nan)
    std_ = data.std()
    mean_ = data.mean()
    if inclusive:
        data[data>mean_+std_*scale]=mean_+std_*scale
        data[data<mean_-std_*scale]=mean_-std_*scale
    else:
        data[data>mean_+std_*scale]=np.nan
        data[data<mean_-std_*scale]=np.nan
    return data

def standardlize(data, inf2nan=True):
    """
    标准化
    :param data:待处理数据
    :param inf2nan：是否将inf转化为nan
    """
    if inf2nan:
        data = data.replace([np.inf, -np.inf], np.nan)
    return (data - data.mean()) / data.std()
    
def neutralize_MarketValue(data, date):
    """
    市值中性化(data以code*date的形式)
    :param data：待处理数据
    :param date：目标日期
    """
    if isinstance(data,pd.Series):
        data = data.to_frame()
    data = data.dropna(how='any')
    security = data.index.to_list()
    market_value = stk_get_daily_mktvalue_pt(symbols=security, fields='tot_mv', trade_date=date, df=True)
    if len(market_value)==0:
        return data
    market_value = market_value[['symbol','tot_mv']].set_index('symbol').dropna(how='any')
    x = sm.add_constant(market_value)
    common_index = list(set(x.index) & set(data.index))
    x = x.loc[common_index,:]
    data = data.loc[common_index,:]
    residual = sm.OLS(data, x).fit().resid# 此处使用最小二乘回归计算
    return residual

def get_N_trading_date(context,date,counts=1,model='previous',exchange='SHSE'):
    """
    获取end_date前N个交易日,end_date为datetime格式，不包括date日期
    :param date：目标日期
    :param counts：历史回溯天数，默认为1，即前一天，该值需为非零正整数
    :param model: 模式，默认为previous，即前N个交易日；next模式为后N个交易日
    """
    if counts<=0 or not isinstance(counts, int):
        print('参数类型错误：get_N_trading_date()的counts参数需为非零正整数.')
        import sys
        sys.exit(-1)
    date = pd.Timestamp(date)
    date_str = date.strftime('%Y-%m-%d')
    # 判断context.trading_dates是否存在
    if 'trading_dates' not in {k: v for k, v in context.__dict__.items() if not k.startswith('__')}.keys():
        context.trading_dates = get_trading_dates_by_year(exchange=exchange, start_year=date.year-1, end_year=datetime.datetime.now().year)
    # 判断当前日期是否为交易日
    pre_times = 1
    if date_str not in list(context.trading_dates['trade_date']) and str(date.year) in set([x[:4] for x in context.trading_dates['trade_date']]):# 非交易日
        date = context.trading_dates[context.trading_dates['date']==date_str]['pre_trade_date'].iloc[0]
        pre_times = 0
    date = pd.Timestamp(date)
    date_str = date.strftime('%Y-%m-%d')
    # 历史N个交易日
    for i in range(2):
        trading_dates_dropna = context.trading_dates.replace('',np.nan).dropna(how='any')
        if model == 'previous':
            previous_date_df = trading_dates_dropna[trading_dates_dropna['trade_date']<=date_str]
            if len(previous_date_df)>=counts+1:
                return previous_date_df['trade_date'].iloc[-counts-pre_times]
            else:
                start_date = pd.Timestamp(context.trading_dates['date'].iloc[0])-datetime.timedelta(days = int(365*np.ceil((counts-len(previous_date_df))/250)))
                start_year = min(start_date.year-1,date.year)
                end_year = int(context.trading_dates['date'].iloc[-1][:4])
                context.trading_dates = get_trading_dates_by_year(exchange=exchange, start_year=start_year, end_year=end_year)
        elif model == 'next':
            previous_date_df = trading_dates_dropna[trading_dates_dropna['trade_date']>=date_str]
            if len(previous_date_df)>=counts+1:
                return previous_date_df['trade_date'].iloc[counts]
            else:
                end_date = pd.Timestamp(context.trading_dates['date'].iloc[-1])+datetime.timedelta(days = int(365*np.ceil((counts-len(previous_date_df))/250)))
                start_year = int(context.trading_dates['date'].iloc[0][:4])
                end_year = max(end_date.year,date.year)
                context.trading_dates = get_trading_dates_by_year(exchange=exchange, start_year=start_year, end_year=end_year)


def get_normal_stocks(date,new_days=365,skip_suspended=True, skip_st=True):
    """
    获取目标日期date的A股代码（剔除停牌股、ST股、次新股（365天））
    :param date：目标日期
    :param new_days:新股上市天数，默认为365天
    """
    date = pd.Timestamp(date).replace(tzinfo=None)
    # A股，剔除停牌和ST股票
    stocks_info = get_symbols(sec_type1=1010, sec_type2=101001, skip_suspended=skip_suspended, skip_st=skip_st, trade_date=date.strftime('%Y-%m-%d'), df=True)
    stocks_info['listed_date'] = stocks_info['listed_date'].apply(lambda x:x.replace(tzinfo=None))
    stocks_info['delisted_date'] = stocks_info['delisted_date'].apply(lambda x:x.replace(tzinfo=None))
    # 剔除次新股和退市股
    stocks_info = stocks_info[(stocks_info['listed_date']<=date-datetime.timedelta(days=new_days))&(stocks_info['delisted_date']>date)]
    all_stocks = list(stocks_info['symbol'])
    # all_stocks = [code for code in all_stocks if code[:8]!='SHSE.688']
    all_stocks_str = ','.join(all_stocks)
    return all_stocks,all_stocks_str


def history_new(security,frequency,start_time,end_time,fields,skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True, benchmark='SHSE.000300'):
    """
    分区间获取数据（以避免超出数据限制）(start_time和end_date为字符串,fields需包含eob和symbol,单字段)
    :param ：参数同history()参数一致，adjust_end_time默认为回测结束时间：None,注意需要根据不同场景使用end_time或context.backtest_end_time
    :param type：默认为True，输出2维DataFrame（日期*股票）,否则输出1维DataFrame
    """
    Data = pd.DataFrame()
    if frequency=='1d':
        trading_date = pd.Series(get_trading_dates(exchange='SZSE', start_date=start_time, end_date=end_time))
    elif frequency=='tick':
        trading_date = history(benchmark, frequency=frequency, start_time=start_time, end_time=end_time, fields='created_at', skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
    else:
        trading_date = history(benchmark, frequency=frequency, start_time=start_time, end_time=end_time, fields='bob,eob', skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
    # 计算合理间隔
    if isinstance(security,str):
        security = security.split(',')
    space = 30000//len(security)
    # 获取数据
    if len(trading_date)<=space:
        Data = history(security, frequency=frequency, start_time=start_time, end_time=end_time, fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
    else:
        for n in range(int(np.ceil(len(trading_date)/space))):
            start = n*space
            end = start+space
            if end>=len(trading_date):
                if frequency=='1d':
                    data = history(security, frequency=frequency, start_time=trading_date.iloc[start], end_time=trading_date.iloc[-1], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
                elif frequency=='tick':
                    data = history(security, frequency=frequency, start_time=trading_date.iloc[start][0], end_time=trading_date.iloc[-1][0], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
                else:
                    data = history(security, frequency=frequency, start_time=trading_date.iloc[start][0], end_time=trading_date.iloc[-1][1], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
            else:
                if frequency=='1d':
                    data = history(security, frequency=frequency, start_time=trading_date.iloc[start], end_time=trading_date.iloc[end], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
                else:
                    data = history(security, frequency=frequency, start_time=trading_date.iloc[start][0], end_time=trading_date.iloc[end][0], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
            if len(data)==33000:
                print('请检查返回数据量，可能超过系统限制，缺少数据！！！！！！！！！！')
            Data = pd.concat([Data,data])
    if frequency=='tick': 
        Data.sort_values(['symbol','created_at'],inplace=True)
        Data.drop_duplicates(subset=['created_at','symbol'],keep='first',inplace=True)
    else:
        Data.sort_values(['symbol','eob'],inplace=True)
        Data.drop_duplicates(subset=['eob','symbol'],keep='first',inplace=True)
    if type:
        if len(Data)>0:
            if frequency=='tick':
                Data = Data.set_index(['created_at','symbol'])
            else:
                Data = Data.set_index(['eob','symbol'])
            Data = Data.unstack()
            Data.columns = Data.columns.droplevel(level=0)
    return Data


def on_order_status(context, order):
    # 标的代码
    symbol = order['symbol']
    # 委托价格
    price = order['price']
    # 委托数量
    volume = order['volume']
    # 目标仓位
    target_percent = order['target_percent']
    # 查看下单后的委托状态，等于3代表委托全部成交
    status = order['status']
    # 买卖方向，1为买入，2为卖出
    side = order['side']
    # 开平仓类型，1为开仓，2为平仓
    effect = order['position_effect']
    # 委托类型，1为限价委托，2为市价委托
    order_type = order['order_type']
    if status == 3:
        if effect == 1:
            if side == 1:
                side_effect = '开多仓'
            elif side == 2:
                side_effect = '开空仓'
        else:
            if side == 1:
                side_effect = '平空仓'
            elif side == 2:
                side_effect = '平多仓'
        order_type_word = '限价' if order_type==1 else '市价'
        print('{}:标的：{}，操作：以{}{}，委托价格：{}，委托数量：{}'.format(context.now,symbol,order_type_word,side_effect,price,volume))
    elif status == 5:
        print('{}:拒单原因：{}'.format(context.now,order['ord_rej_reason_detail']))
       

def on_backtest_finished(context, indicator):
    print('*'*50)
    print('回测已完成，请通过右上角“回测历史”功能查询详情。')
    

if __name__ == '__main__':
    '''
        strategy_id策略ID, 由系统生成
        filename文件名, 请与本文件名保持一致
        mode运行模式, 实时模式:MODE_LIVE回测模式:MODE_BACKTEST
        token绑定计算机的ID, 可在系统设置-密钥管理中生成
        backtest_start_time回测开始时间
        backtest_end_time回测结束时间
        backtest_adjust股票复权方式, 不复权:ADJUST_NONE前复权:ADJUST_PREV后复权:ADJUST_POST
        backtest_initial_cash回测初始资金
        backtest_commission_ratio回测佣金比例
        backtest_slippage_ratio回测滑点比例
        '''
    run(strategy_id='a51d4ab3-462b-11ee-b0d2-f46b8c02346f',
        filename='main.py',
        mode=MODE_BACKTEST,
        token='47ca47f849b3a0f66ec0f7013bb56bb667d63a70',
        backtest_start_time='2018-12-28 08:00:00',
        backtest_end_time='2023-10-08 16:00:00',
        backtest_adjust=ADJUST_PREV,
        backtest_initial_cash=10000000,
        backtest_commission_ratio=0.0007,# 买入万二手续费+卖出万二手续费和千1印花税，免5
        backtest_slippage_ratio=0.00123)